{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cb83dcd9-c16a-4d4a-99de-3770b0f0fd33",
   "metadata": {},
   "outputs": [],
   "source": [
    "#导包\n",
    "%matplotlib inline\n",
    "\n",
    "import math\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "import dltools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db5d87e8-82b7-4eb6-8453-756b8b377e3b",
   "metadata": {},
   "source": [
    "### 加载time machine数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4d0c2ca2-1884-4690-96ae-0ac2c6ae5aac",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size, num_steps = 32, 15\n",
    "# 默认需要从亚马逊aws云上面下载数据, 但是资源过期了, 下载不了了. 导致数据加载不了. \n",
    "train_iter, vocab = dltools.load_data_time_machine(batch_size=batch_size, num_steps=num_steps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "92d9929d-ee1c-4680-8d34-fea218a8b978",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[ 2,  1,  3,  5, 13,  2,  1, 13,  4, 15,  9,  5,  6,  2,  1],\n",
      "        [14, 18,  9,  3,  1,  3,  9,  2,  1, 21, 14, 21, 21, 12,  2],\n",
      "        [ 3,  9,  1,  4,  1, 12,  2,  4,  6,  1, 16,  7, 10,  2, 16],\n",
      "        [ 6,  7,  3,  1,  3,  9,  4,  3,  1, 10,  4,  3,  9,  2, 10],\n",
      "        [15,  2,  1,  3,  9,  2, 19,  1,  3,  4, 14, 18,  9,  3,  1],\n",
      "        [13,  2,  6,  3,  1, 15,  4,  6,  1,  4,  6,  1,  5,  6,  8],\n",
      "        [12,  5,  6, 16,  5, 10, 13,  5,  3, 19,  1,  7, 16,  1,  3],\n",
      "        [20, 20,  2,  6,  8,  1,  3,  9,  4,  3,  7, 14, 10,  1, 15],\n",
      "        [13,  2,  1,  3, 10,  4, 22,  2, 12, 12,  2, 10,  1, 17,  5],\n",
      "        [ 7, 14, 10,  1, 15,  7,  6,  8, 15,  5,  7, 14,  8,  6,  2],\n",
      "        [13,  4, 19,  1, 15,  4, 12, 12,  1, 12,  2,  6, 18,  3,  9],\n",
      "        [ 1, 18,  2,  7, 13,  2,  3, 10, 19, 20, 10,  7, 16,  2,  8],\n",
      "        [ 1, 16,  7, 14, 10,  1,  5, 16,  1,  3,  9,  2, 19,  1, 15],\n",
      "        [ 1,  3,  2, 12, 12,  5,  6, 18,  1, 19,  7, 14,  1,  5,  1],\n",
      "        [ 8,  2,  6,  3,  4,  3,  5,  7,  6,  8,  1,  7, 16,  1,  9],\n",
      "        [16,  5,  6, 18,  2, 10,  1,  8,  9,  7, 17,  8,  1,  3,  9],\n",
      "        [ 3,  1, 15,  7,  6, 15, 12, 14, 11,  2,  1, 17,  4,  8,  1],\n",
      "        [10,  1,  8, 13,  5, 12,  2, 11,  1,  4, 10,  2,  1, 19,  7],\n",
      "        [ 1,  8, 20,  4,  8, 13,  7, 11,  5, 15,  1, 26, 14, 13, 20],\n",
      "        [ 5,  8,  1, 26, 14,  8,  3,  1, 17,  9,  2, 10,  2,  3,  9],\n",
      "        [ 2,  8,  1,  4, 21,  7, 22,  2,  1,  3,  9,  2,  1,  2,  4],\n",
      "        [ 5, 22,  5, 11, 12, 19,  1,  5,  1, 18,  7,  1, 21,  4, 15],\n",
      "        [ 1,  9,  2,  1, 15,  4,  6,  1, 18,  7,  1, 14, 20,  1,  4],\n",
      "        [13,  2,  1,  3, 10,  4, 22,  2, 12, 12,  2, 10, 19,  7, 14],\n",
      "        [13,  4,  6,  3,  9,  4,  3,  1,  8,  9,  4, 12, 12,  1,  3],\n",
      "        [ 1,  4, 15, 15,  7, 14,  6,  3,  1,  7, 16,  1,  3,  9,  2],\n",
      "        [ 2,  1, 18,  7,  3,  9,  2,  1, 18,  2, 10, 13,  4,  6,  1],\n",
      "        [15,  9,  7, 12,  7, 18,  5,  8,  3, 19,  2,  8,  1,  8,  7],\n",
      "        [12,  1,  8, 13,  5, 12,  5,  6, 18,  1, 16,  4,  5,  6,  3],\n",
      "        [ 4,  1, 15,  7,  6, 26, 14, 10,  2, 10,  1,  9,  2,  1,  9],\n",
      "        [15, 10, 19,  8,  3,  4, 12, 12,  5,  6,  2,  1,  8, 14, 21],\n",
      "        [ 1, 13,  2, 15,  9,  4,  6,  5,  8, 13,  3,  9,  2,  6,  1]]) tensor([[ 1,  3,  5, 13,  2,  1, 13,  4, 15,  9,  5,  6,  2,  1, 21],\n",
      "        [18,  9,  3,  1,  3,  9,  2,  1, 21, 14, 21, 21, 12,  2,  8],\n",
      "        [ 9,  1,  4,  1, 12,  2,  4,  6,  1, 16,  7, 10,  2, 16,  5],\n",
      "        [ 7,  3,  1,  3,  9,  4,  3,  1, 10,  4,  3,  9,  2, 10,  1],\n",
      "        [ 2,  1,  3,  9,  2, 19,  1,  3,  4, 14, 18,  9,  3,  1, 19],\n",
      "        [ 2,  6,  3,  1, 15,  4,  6,  1,  4,  6,  1,  5,  6,  8,  3],\n",
      "        [ 5,  6, 16,  5, 10, 13,  5,  3, 19,  1,  7, 16,  1,  3,  9],\n",
      "        [20,  2,  6,  8,  1,  3,  9,  4,  3,  7, 14, 10,  1, 15,  7],\n",
      "        [ 2,  1,  3, 10,  4, 22,  2, 12, 12,  2, 10,  1, 17,  5,  3],\n",
      "        [14, 10,  1, 15,  7,  6,  8, 15,  5,  7, 14,  8,  6,  2,  8],\n",
      "        [ 4, 19,  1, 15,  4, 12, 12,  1, 12,  2,  6, 18,  3,  9, 21],\n",
      "        [18,  2,  7, 13,  2,  3, 10, 19, 20, 10,  7, 16,  2,  8,  8],\n",
      "        [16,  7, 14, 10,  1,  5, 16,  1,  3,  9,  2, 19,  1, 15,  7],\n",
      "        [ 3,  2, 12, 12,  5,  6, 18,  1, 19,  7, 14,  1,  5,  1,  9],\n",
      "        [ 2,  6,  3,  4,  3,  5,  7,  6,  8,  1,  7, 16,  1,  9,  5],\n",
      "        [ 5,  6, 18,  2, 10,  1,  8,  9,  7, 17,  8,  1,  3,  9,  2],\n",
      "        [ 1, 15,  7,  6, 15, 12, 14, 11,  2,  1, 17,  4,  8,  1,  4],\n",
      "        [ 1,  8, 13,  5, 12,  2, 11,  1,  4, 10,  2,  1, 19,  7, 14],\n",
      "        [ 8, 20,  4,  8, 13,  7, 11,  5, 15,  1, 26, 14, 13, 20,  5],\n",
      "        [ 8,  1, 26, 14,  8,  3,  1, 17,  9,  2, 10,  2,  3,  9,  2],\n",
      "        [ 8,  1,  4, 21,  7, 22,  2,  1,  3,  9,  2,  1,  2,  4, 10],\n",
      "        [22,  5, 11, 12, 19,  1,  5,  1, 18,  7,  1, 21,  4, 15, 23],\n",
      "        [ 9,  2,  1, 15,  4,  6,  1, 18,  7,  1, 14, 20,  1,  4, 18],\n",
      "        [ 2,  1,  3, 10,  4, 22,  2, 12, 12,  2, 10, 19,  7, 14,  1],\n",
      "        [ 4,  6,  3,  9,  4,  3,  1,  8,  9,  4, 12, 12,  1,  3, 10],\n",
      "        [ 4, 15, 15,  7, 14,  6,  3,  1,  7, 16,  1,  3,  9,  2,  1],\n",
      "        [ 1, 18,  7,  3,  9,  2,  1, 18,  2, 10, 13,  4,  6,  1,  8],\n",
      "        [ 9,  7, 12,  7, 18,  5,  8,  3, 19,  2,  8,  1,  8,  7,  1],\n",
      "        [ 1,  8, 13,  5, 12,  5,  6, 18,  1, 16,  4,  5,  6,  3, 12],\n",
      "        [ 1, 15,  7,  6, 26, 14, 10,  2, 10,  1,  9,  2,  1,  9,  4],\n",
      "        [10, 19,  8,  3,  4, 12, 12,  5,  6,  2,  1,  8, 14, 21,  8],\n",
      "        [13,  2, 15,  9,  4,  6,  5,  8, 13,  3,  9,  2,  6,  1,  9]])\n"
     ]
    }
   ],
   "source": [
    "for x, y in train_iter:\n",
    "    print(x, y)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "3bd4cd76-575e-41d7-83a7-54dbbbb94729",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 15])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "73172ec7-19c4-4205-b990-25bd1da7d7e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 15])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ece4cb47-2c44-464f-8a42-a92b4e640dd4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 1, 18,  9,  7,  2,  2,  5, 20,  2, 14,  4, 18, 16,  3,  2,  5,  1,  1,\n",
       "         8,  8,  8, 22,  9,  2,  4,  4,  1,  9,  1,  1, 10, 13,  3,  9,  1,  3,\n",
       "         1,  6,  6,  2,  1, 10, 19,  2,  7,  2,  6,  6, 15,  8, 20,  1,  1,  5,\n",
       "         2,  1,  6, 15, 18,  7,  8, 15, 19,  2,  5,  3,  4,  1,  3,  3, 16,  6,\n",
       "         3,  1,  1,  7, 14, 12,  3, 18,  7, 13,  4, 26,  4, 11,  1,  3,  3, 15,\n",
       "         7, 12, 13,  7,  8, 15, 13,  1,  1,  3,  9,  1,  5,  8, 10, 15, 15, 13,\n",
       "        10, 12,  4,  2,  6,  5,  8, 14, 21, 12, 15, 10,  9,  7,  3,  7,  5,  6,\n",
       "         3,  9,  2,  3, 12,  9,  2, 15, 10,  1,  4,  7,  4,  2,  1,  5,  3, 10,\n",
       "        15, 12, 13,  8,  7, 19,  4,  4,  4, 14,  9, 18, 12, 26,  4,  4,  1,  9,\n",
       "         2,  4, 19,  4, 13,  3, 22,  6, 12,  3,  5,  6,  5,  1, 12,  2,  7,  3,\n",
       "        22,  1,  6, 22,  3,  6,  2,  5,  5, 14, 12,  6, 13,  2,  4,  3,  1,  6,\n",
       "         5,  9,  2,  8, 12, 10, 16, 18,  7,  8, 14, 11, 11,  1,  2,  5,  1,  2,\n",
       "         1,  3,  1,  8,  6, 10, 12,  5,  4,  1,  6,  1,  3,  1,  3,  4, 12, 15,\n",
       "         1, 19,  1,  1,  6,  9, 11,  1,  5, 17,  1,  1, 18, 12,  8,  1, 18,  3,\n",
       "        18,  2,  5,  8, 15, 21,  1, 10,  4,  4, 19,  3, 12,  5, 12, 20,  3, 19,\n",
       "         8,  7,  2,  4, 15,  9,  3, 18,  7, 12,  9,  7,  2, 19,  1, 10,  6, 13,\n",
       "         9, 14, 16,  4, 14,  6,  1,  7,  2,  7,  2, 10,  9,  7,  1, 17,  1, 10,\n",
       "         1,  2,  9,  7,  1,  2,  4, 16, 10,  2, 16,  1,  2,  3,  5, 21,  7,  3,\n",
       "        18,  1,  7, 14, 10, 14,  6,  7,  2, 14,  7,  8, 17,  2, 26, 10,  2,  1,\n",
       "        14, 10, 12,  1, 13,  8,  4,  9,  1,  9,  6, 21, 10,  9,  9,  5, 16, 10,\n",
       "         1,  8, 18, 16, 19,  1, 16,  1,  4,  1, 14,  2,  1, 21, 20, 19, 12,  3,\n",
       "         4,  1,  5,  2,  8,  2,  2, 12,  2,  2,  3,  6,  1,  1, 17,  6,  3,  2,\n",
       "         1,  5,  1,  3,  8, 19, 13,  3,  2,  4,  1,  7,  1,  9,  6,  8,  6,  1,\n",
       "        14,  6,  1,  2, 16, 10,  1,  8,  3, 15,  5,  2,  9,  8, 15,  1,  9,  9,\n",
       "         1,  7, 20,  9,  4, 15,  4, 14,  3,  2,  1,  7,  3,  9, 21,  1, 21,  8,\n",
       "         5,  1, 19,  3,  9,  7,  3,  8, 21,  8,  7,  9,  5,  2,  4, 14,  5,  2,\n",
       "        10, 23, 18,  1, 10,  1,  8,  1, 12,  4,  8,  9])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y.T.reshape(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "519caa64-eade-4c37-9334-82ddced175a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['<unk>',\n",
       " ' ',\n",
       " 'e',\n",
       " 't',\n",
       " 'a',\n",
       " 'i',\n",
       " 'n',\n",
       " 'o',\n",
       " 's',\n",
       " 'h',\n",
       " 'r',\n",
       " 'd',\n",
       " 'l',\n",
       " 'm',\n",
       " 'u',\n",
       " 'c',\n",
       " 'f',\n",
       " 'w',\n",
       " 'g',\n",
       " 'y',\n",
       " 'p',\n",
       " 'b',\n",
       " 'v',\n",
       " 'k',\n",
       " 'x',\n",
       " 'z',\n",
       " 'j',\n",
       " 'q']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.idx_to_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "57282d27-a96c-46ab-80c5-f40f435e6cf8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'<unk>': 0,\n",
       " ' ': 1,\n",
       " 'e': 2,\n",
       " 't': 3,\n",
       " 'a': 4,\n",
       " 'i': 5,\n",
       " 'n': 6,\n",
       " 'o': 7,\n",
       " 's': 8,\n",
       " 'h': 9,\n",
       " 'r': 10,\n",
       " 'd': 11,\n",
       " 'l': 12,\n",
       " 'm': 13,\n",
       " 'u': 14,\n",
       " 'c': 15,\n",
       " 'f': 16,\n",
       " 'w': 17,\n",
       " 'g': 18,\n",
       " 'y': 19,\n",
       " 'p': 20,\n",
       " 'b': 21,\n",
       " 'v': 22,\n",
       " 'k': 23,\n",
       " 'x': 24,\n",
       " 'z': 25,\n",
       " 'j': 26,\n",
       " 'q': 27}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab.token_to_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "9ea25226-4fb5-4261-b284-102b614838bb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "28"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d798af4b-81f9-487c-b34e-9722ea8df6ab",
   "metadata": {},
   "source": [
    "### one_hot编码的数据输入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "212b6333-718f-4205-a72f-ffe6ae3e2e49",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0],\n",
       "        [0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "         0, 0, 0, 0]])"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# pytorch提供了快速进行one_hot编码的工具\n",
    "F.one_hot(torch.tensor([0,2]),  num_classes=len(vocab))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7de8e9d0-ebcd-48cd-bc92-d95b54880d50",
   "metadata": {},
   "outputs": [],
   "source": [
    "X = torch.arange(10).reshape((2,5))\n",
    "x = F.one_hot(X.T, 28)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "9c781adf-35f5-47a4-9a49-51de430c999e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 2, 28])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ae57ff6e-cc79-4aaf-869a-d564ed3ffdb5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 5],\n",
       "        [1, 6],\n",
       "        [2, 7],\n",
       "        [3, 8],\n",
       "        [4, 9]])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "78ec9019-e24f-4719-a5d5-15300d97e1e9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0]],\n",
       "\n",
       "        [[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0]],\n",
       "\n",
       "        [[0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0]],\n",
       "\n",
       "        [[0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0]],\n",
       "\n",
       "        [[0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0],\n",
       "         [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,\n",
       "          0, 0, 0, 0, 0]]])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "3cc456a9-9120-4ac0-89c7-1253176de532",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 28])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x[0].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d8414a2-54f9-4663-ac42-e58d1a3e0a58",
   "metadata": {},
   "source": [
    "### 初始化模型参数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "6b1f04c5-425e-49c1-b716-8d4f78d3d9cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_params(vocab_size, num_hiddens, device):\n",
    "    num_inputs = num_outputs = vocab_size\n",
    "    def normal(shape):\n",
    "        return torch.randn(size=shape, device=device) * 0.01\n",
    "    # 隐藏层参数\n",
    "    W_xh = normal((num_inputs, num_hiddens))\n",
    "    W_hh = normal((num_hiddens, num_hiddens))\n",
    "    b_h = torch.zeros(num_outputs, device=device)\n",
    "     # 输出层参数\n",
    "    W_hq = normal((num_hiddens, num_outputs))\n",
    "    b_q = torch.zeros(num_outputs, device=device)\n",
    "    # 把这些参数都设置requires_grad = True\n",
    "    params = [W_xh, W_hh, b_h, W_hq, b_q]\n",
    "    for param in params:\n",
    "        param.requires_grad_(True)\n",
    "    return params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "5d5c9968-2a10-41eb-9bae-30c13bccb2f6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([[-0.0089,  0.0074,  0.0114,  ...,  0.0063, -0.0047, -0.0070],\n",
       "         [-0.0055, -0.0061, -0.0182,  ..., -0.0076,  0.0061, -0.0005],\n",
       "         [-0.0016,  0.0020, -0.0047,  ...,  0.0054,  0.0142, -0.0038],\n",
       "         ...,\n",
       "         [-0.0021, -0.0141,  0.0036,  ...,  0.0141,  0.0057, -0.0008],\n",
       "         [-0.0191, -0.0224, -0.0046,  ...,  0.0134,  0.0009,  0.0102],\n",
       "         [ 0.0128,  0.0129, -0.0064,  ...,  0.0067, -0.0123, -0.0022]],\n",
       "        requires_grad=True),\n",
       " tensor([[-0.0070, -0.0156,  0.0036,  ..., -0.0071,  0.0053, -0.0035],\n",
       "         [ 0.0061, -0.0075, -0.0140,  ...,  0.0109,  0.0046,  0.0176],\n",
       "         [-0.0070,  0.0106, -0.0065,  ...,  0.0131, -0.0050, -0.0023],\n",
       "         ...,\n",
       "         [ 0.0021,  0.0071,  0.0006,  ...,  0.0115, -0.0015, -0.0043],\n",
       "         [ 0.0058,  0.0138, -0.0122,  ...,  0.0248,  0.0062,  0.0085],\n",
       "         [-0.0033,  0.0036, -0.0145,  ..., -0.0009,  0.0131, -0.0240]],\n",
       "        requires_grad=True),\n",
       " tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0.], requires_grad=True),\n",
       " tensor([[-0.0035,  0.0101,  0.0158,  ...,  0.0002, -0.0010, -0.0075],\n",
       "         [-0.0094, -0.0129, -0.0023,  ..., -0.0229,  0.0126, -0.0050],\n",
       "         [ 0.0002,  0.0130, -0.0092,  ...,  0.0249, -0.0046,  0.0010],\n",
       "         ...,\n",
       "         [ 0.0065, -0.0084, -0.0099,  ..., -0.0054,  0.0030,  0.0046],\n",
       "         [-0.0206, -0.0031, -0.0100,  ...,  0.0087,  0.0089, -0.0080],\n",
       "         [ 0.0057, -0.0116, -0.0114,  ...,  0.0052, -0.0093, -0.0079]],\n",
       "        requires_grad=True),\n",
       " tensor([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,\n",
       "         0., 0., 0., 0.], requires_grad=True)]"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "get_params(28,512,'cpu:0')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af626379-4bd0-4c1f-9f78-c7231cabb9e2",
   "metadata": {},
   "source": [
    "### 初始化时返回隐藏状态"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "f38c42aa-a225-40c6-930b-bf16aa24d2f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_rnn_state(batch_size, num_hiddens, device):\n",
    "    # 返回的是一个元组\n",
    "    return (torch.zeros((batch_size, num_hiddens), device=device),)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b89c50-0c77-42c2-bc59-5d67babf8a90",
   "metadata": {},
   "source": [
    "### rnn主体结构"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "0751d332-9865-4a58-8fef-5bde293104f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rnn(inputs, state, params):\n",
    "    #inputs的形状: (时间步数量, 批次大小, 词表大小)\n",
    "    W_xh, W_hh, b_h, W_hq, b_q = params\n",
    "    H, = state\n",
    "    outputs = []\n",
    "    # X的shape: [批次大小, 词表大小]\n",
    "    for X in inputs:\n",
    "        # 一般在循环神经网络中激活函数用tanh\n",
    "        H = torch.tanh(torch.mm(X, W_xh) + torch.mm(H, W_hh) + b_h)\n",
    "        Y = torch.mm(H, W_hq) + b_q\n",
    "        outputs.append(Y)\n",
    "    return torch.cat(outputs, dim=0),(H,)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "22b6f9c9-4769-4d51-aef8-4d7d7b8bf4e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 包装成类\n",
    "class RNNModelScratch:\n",
    "    def __init__(self, vocab_size, num_hiddens, device, get_params, init_state, forward_fn):\n",
    "        self.vocab_size, self.num_hiddens = vocab_size, num_hiddens\n",
    "        self.params = get_params(vocab_size, num_hiddens, device)\n",
    "        self.init_state, self.forward_fn = init_state, forward_fn\n",
    "        \n",
    "    def __call__(self, X, state):\n",
    "        X = F.one_hot(X.T, self.vocab_size).type(torch.float32)\n",
    "        return self.forward_fn(X, state, self.params)\n",
    "    \n",
    "    def begin_state(self, batch_size, device):\n",
    "        return self.init_state(batch_size, self.num_hiddens, device)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "124f96c0-85f5-48bb-a403-69423fd42114",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "7e5ede22-8e64-4b1e-80f5-98e63d349d76",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cpu'"
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "1a86232c-e112-4ea7-86a2-6be1d19e954f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dltools.try_gpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "43c7bd38-8b39-46f4-818b-b41984c24955",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 5])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "0e56cf15-e05b-4f76-983d-5c4ac82eb852",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 试用一下. \n",
    "num_hiddens = 28\n",
    "net = RNNModelScratch(len(vocab), num_hiddens, dltools.try_gpu(), get_params, init_rnn_state, rnn)\n",
    "state = net.begin_state(X.shape[0], dltools.try_gpu())\n",
    "Y, new_state = net(X.to(dltools.try_gpu()), state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "d3e112db-46ce-4fc6-aae5-9ee11559e93e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([10, 28])"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Y.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "1c39ef39-df98-40f7-9a14-78f2a523ee26",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(new_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "633028b6-37c7-4190-b98c-b6de8e42a4b0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([2, 28])"
      ]
     },
     "execution_count": 48,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_state[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "fedf7ad8-9258-431c-a693-f462bbbaf06d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vocab['a']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "66f2b73b-6998-4d7c-9e1d-9ccba0ceb09c",
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = [4,3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8615ec4f-134d-4ed2-9a53-1c4695293e27",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
